import cv2
import numpy as np
from chinese_font import cv2AddChineseText


def distinguish_face_camera(trainer_file, predicted_file):
    predicted_person = ''
    # 视频播放器
    # video = "http://admin:224114@192.168.137.212:8081/video"
    # capture = cv2.VideoCapture(video)
    capture = cv2.VideoCapture(0)
    # LBPH人脸识别器
    recognizer = cv2.face.LBPHFaceRecognizer_create()
    # LBPH读取训练好的模型
    recognizer.read(trainer_file)
    predicted_file = predicted_file
    # DNN人脸检测
    net = cv2.dnn.readNetFromTensorflow("dataset/DNN_model/opencv_face_detector_uint8.pb",
                                        "dataset/DNN_model/opencv_face_detector.pbtxt")
    # 读取预测文件
    with open(predicted_file, 'r') as f:
        predicted_names = f.readlines()
    # 摄像头没有打开
    if not capture.isOpened():
        print("视频没有打开")
        return False
    # 循环播放
    while True:
        ret, frame = capture.read()
        # 播放完毕
        if not ret:
            print("播放完毕")
            return False
        # 调整视频大小
        if frame.shape[0] < 500 and frame.shape[1] < 500:
            frame = cv2.resize(frame, (frame.shape[1] * 2, frame.shape[0] * 2))
        elif frame.shape[0] > 1000 or frame.shape[1] > 1000:
            frame = cv2.resize(frame, (frame.shape[1] // 2, frame.shape[0] // 2))
        #
        if ret:
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            (h, w) = gray.shape[:2]
            # 人脸检测
            blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300), [104, 177, 123], False, False)
            net.setInput(blob)
            detections = net.forward()
            for i in range(detections.shape[2]):
                confidence = detections[0, 0, i, 2]
                if confidence > 0.7:
                    box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                    (startX, startY, endX, endY) = box.astype('int')
                    # 模型预测
                    predicted_index, conf = recognizer.predict(gray[startY:endY, startX:endX])
                    # 预测结果同预测文件对比，得出人脸对象名词
                    for j in range(len(predicted_names)):
                        if str(predicted_index) == str(predicted_names[j].split(':')[0]):
                            predicted_person = str(predicted_names[j].split(':')[1])
                            break
                    # 图像标记
                    cv2.rectangle(frame, (startX, startY), (endX, endY), (255, 0, 0), 3)
                    frame = cv2AddChineseText(frame, str(predicted_person), (startX + 20, startY + 20),
                                              (255, 255, 0), 30)
            # 图像显示
            cv2.imshow('test', frame)
        # 退出 q
        if cv2.waitKey(28) & 0xFF == ord('q'):
            break
    # 关闭播放器
    capture.release()
    cv2.destroyAllWindows()


# trainer_file_path = "dataset/trainer.yam"
# predicted_file_path = "dataset/predicted.txt"
# distinguish_face_camera(trainer_file_path, predicted_file_path)
